POC Memory Tracking of Potentially Shared Buffer #6590

tustvold · 2024-10-19T09:20:00Z

Which issue does this PR close?

Relates to #6439.

Rationale for this change

A POC showing how we could add more accurate memory tracking to arrays without a huge amount of churn

What changes are included in this PR?

This adds two traits MemoryPool and MemoryReservation that work together to allow accurate instrumentation of Buffer memory usage.

Are there any user-facing changes?

tustvold

I think this should provide a mechanism for consumers like DF to track memory usage of buffered arrays accurately, and with performance that is no worse than the current get_buffer_memory_size, which this would likely lead us to deprecate.

tustvold · 2024-10-19T09:21:02Z

arrow-buffer/src/bytes.rs

@@ -44,6 +44,9 @@ pub struct Bytes {

    /// how to deallocate this region
    deallocation: Deallocation,
+
+    #[cfg(feature = "pool")]
+    reservation: std::sync::Mutex<Option<Box<dyn crate::MemoryReservation>>>,


I originally used parking_lot::Mutex but this would have made Buffer not RewindSafe, which would be a breaking change (and given there is a test for this I presume this is important).

tustvold · 2024-10-19T09:22:06Z

arrow-buffer/src/pool.rs

+
+/// A memory reservation within a [`MemoryPool`] that is freed on drop
+pub trait MemoryReservation {
+    fn resize(&mut self, new: usize);


I don't actually have a use-case for this atm, but this seemed sensible to add now to avoid it being difficult later

Mutable buffer might need this. It may resize its reservation dynamically

tustvold · 2024-10-19T09:23:10Z

arrow-buffer/src/bytes.rs

+    /// Register this [`Bytes`] with the provided [`MemoryPool`]
+    #[cfg(feature = "pool")]
+    pub fn claim(&self, pool: &dyn crate::MemoryPool) {
+        *self.reservation.lock().unwrap() = Some(pool.register(self.capacity()));


There are two implications of this formulation:

We always allocate a new memory reservation which could be wasteful

We allocate from the new pool before freeing the memory from the previous reservation

tustvold · 2024-10-19T09:23:56Z

arrow-buffer/src/buffer/immutable.rs

+
+    /// Register this [`Buffer`] with the provided [`MemoryPool`]
+    #[cfg(feature = "pool")]
+    pub fn claim(&self, pool: &dyn crate::MemoryPool) {


We could plumb this claim API through the various arrays and RecordBatch. Consumers like DF could then claim RecordBatch they wish to buffer, and know they aren't double-counting.

That is interesting -- so the idea is that this would support cooperatively assigning / tracking usage (rather than relying on some global allocator to do so).

That aligns pretty nicely with the "no overhead unless asked" and "everything being explicit" principles

I think this API would be easier to use / manage if it also returned the prior reservation, if any

I think this would make it harder to chain through to Arrays and RecordBatch that may consist of multiple buffers and therefore reservations. It would also mean potentially keeping around "duplicate" reservations for longer than necessary, which would become even more problematic if the same buffer is used multiple times.

What would be the use-case for the returned reservation?

I was thinking like if only some Arrays in a RecordBatch have a reservation but some didn't and then a system like the DataFusion memory pool added an entirely new reservation for all the arrays wiping out the old reservations

In that case I would want the previous reservations to be unregistered

Although maybe Drop that did the unregistering would be good enough for that case 🤔

Although maybe Drop that did the unregistering would be good enough for that case 🤔

That's what is intended, and follows RAII best-practices

tustvold · 2024-10-19T09:25:36Z

arrow-buffer/Cargo.toml

 [lib]
 name = "arrow_buffer"
 path = "src/lib.rs"
 bench = false

+[features]
+pool = []


As it is very hard to guage what if any performance implications there are of this, I think it is important that it is gated by a feature flag at least initially.

alamb

Thank you @tustvold -- this looks very cool

My suggested next steps are twofold:

Write up some self contained example showing how to use this API in a real world scenario (perhaps @haohuaijin 's The get_array_memory_size() get wrong result(with different compression method) after deconde record from ipc format #6363 of tracking memory sent to/from IPC buffers)
Try and plumb it into DataFusion's MemoryPool and see how that looks / if this API is sufficient

cc @waynexia and @jhorstmann in case you are interested

alamb · 2024-10-21T11:24:51Z

arrow-buffer/src/buffer/immutable.rs

+
+    /// Register this [`Buffer`] with the provided [`MemoryPool`]
+    #[cfg(feature = "pool")]
+    pub fn claim(&self, pool: &dyn crate::MemoryPool) {


That is interesting -- so the idea is that this would support cooperatively assigning / tracking usage (rather than relying on some global allocator to do so).

That aligns pretty nicely with the "no overhead unless asked" and "everything being explicit" principles

alamb · 2024-10-21T11:25:36Z

arrow-buffer/src/buffer/immutable.rs

+
+    /// Register this [`Buffer`] with the provided [`MemoryPool`]
+    #[cfg(feature = "pool")]
+    pub fn claim(&self, pool: &dyn crate::MemoryPool) {


I think this API would be easier to use / manage if it also returned the prior reservation, if any

waynexia · 2024-10-21T14:21:55Z

Thanks for bringing this up 👍 I'll take a closer look tomorrow

waynexia · 2024-10-22T14:54:05Z

This looks fits the requirement of precise memory tracking without touching the unstable rust feature 💯

Adding an explicit API looks fine to me. I'd like to give it a try on mutable buffer either to see if there is any hiding problem

waynexia · 2024-10-25T09:49:01Z

Implementation for MutableBuffer is tustvold#82

On writing that patch I find resize is indeed useful. And in one case I need .resize(0) to reflect one memory is deallocated.

This API LGTM in general.

POC MemoryPool

6646dd4

github-actions bot added the arrow Changes to the arrow crate label Oct 19, 2024

tustvold commented Oct 19, 2024

View reviewed changes

tustvold mentioned this pull request Oct 19, 2024

[DISSCUSSION] the memory tracking approach that better handles shared buffers #6439

Open

tustvold requested a review from alamb October 19, 2024 09:27

Docs

d408327

tustvold changed the title ~~POC MemoryPool~~ POC Memory Tracking of Potentially Shared Buffer Oct 19, 2024

tustvold mentioned this pull request Oct 19, 2024

feat: add Allocator type param to MutableBuffer #6336

Draft

alamb reviewed Oct 21, 2024

View reviewed changes

waynexia self-requested a review October 21, 2024 14:20

waynexia mentioned this pull request Oct 25, 2024

feat: implement memory pool for mutable buffer tustvold/arrow-rs#82

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

POC Memory Tracking of Potentially Shared Buffer #6590

POC Memory Tracking of Potentially Shared Buffer #6590

tustvold commented Oct 19, 2024 •

edited

Loading

tustvold left a comment •

edited

Loading

tustvold Oct 19, 2024

tustvold Oct 19, 2024

waynexia Oct 22, 2024

tustvold Oct 19, 2024 •

edited

Loading

tustvold Oct 19, 2024

alamb Oct 21, 2024

alamb Oct 21, 2024

tustvold Oct 21, 2024 •

edited

Loading

alamb Oct 21, 2024 •

edited

Loading

tustvold Oct 21, 2024

tustvold Oct 19, 2024

alamb left a comment

alamb Oct 21, 2024

alamb Oct 21, 2024

waynexia commented Oct 21, 2024

waynexia commented Oct 22, 2024

waynexia commented Oct 25, 2024

POC Memory Tracking of Potentially Shared Buffer #6590

Are you sure you want to change the base?

POC Memory Tracking of Potentially Shared Buffer #6590

Conversation

tustvold commented Oct 19, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

tustvold left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold Oct 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

alamb Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

waynexia commented Oct 21, 2024

waynexia commented Oct 22, 2024

waynexia commented Oct 25, 2024

tustvold commented Oct 19, 2024 •

edited

Loading

tustvold left a comment •

edited

Loading

tustvold Oct 19, 2024 •

edited

Loading

tustvold Oct 21, 2024 •

edited

Loading

alamb Oct 21, 2024 •

edited

Loading